#include <iostream>
#include <algorithm>
#include <chrono>
#include <stdio.h>
#include <cuda_runtime.h>
#include <cstring>

// 定义常量
#define MAXBIT 32
#define DATASIZE 1000000
#define BLOCK_SIZE 256 // 必须是 2 的幂
// #define DATASIZE 5
// #define BLOCK_SIZE 2 
// #define MAXBIT 32

using namespace std::chrono;

#define CUDA_CHECK(call) \
    do { \
        cudaError_t err = call; \
        if (err != cudaSuccess) { \
            fprintf(stderr, "CUDA error at %s:%d - %s\n", __FILE__, __LINE__, cudaGetErrorString(err)); \
            exit(1); \
        } \
    } while(0)

bool validate(unsigned int *a, unsigned int* b, int n){
    for (int i = 0; i < n; i++) {
        if(a[i] != b[i]){
            printf("index = %d is not matched.\n", i);
            return false;
        }
    }
    return true;
}

void print_array(unsigned int* arr, int n, const char* label) {
    if(n > 20) return;
    printf("\n%s\n", label);
    for (int i = 0; i < n; i++) {
        printf("%u  ", arr[i]);
    }
    printf("\n");
}

void cpu_radix_sort_with_index(unsigned int* arr, unsigned int* arr_pos,
                               unsigned int* sorted_arr, unsigned int* original_indices,
                               int n) {
    unsigned int* values_curr = new unsigned int[n];
    unsigned int* values_next = new unsigned int[n];
    unsigned int* pos_curr = new unsigned int[n];
    unsigned int* pos_next = new unsigned int[n];
    std::copy(arr, arr + n, values_curr);
    std::copy(arr_pos, arr_pos + n, pos_curr);
    for (int bit = 0; bit < MAXBIT; bit++) {
        int count[2] = {0, 0};
        for (int i = 0; i < n; i++) {
            int digit = (values_curr[i] >> bit) & 1;
            count[digit]++;
        }
        int offset[2] = {0, count[0]};
        for (int i = 0; i < n; i++) {
            int digit = (values_curr[i] >> bit) & 1;
            int dest_idx = offset[digit]++;
            values_next[dest_idx] = values_curr[i];
            pos_next[dest_idx] = pos_curr[i];
        }
        std::swap(values_curr, values_next);
        std::swap(pos_curr, pos_next);
    }
    std::copy(values_curr, values_curr + n, sorted_arr);
    std::copy(pos_curr, pos_curr + n, original_indices);

    delete[] values_curr;
    delete[] values_next;
    delete[] pos_curr;
    delete[] pos_next;
}

// 提取当前 bit
__global__ void computeBitValues(unsigned int* d_in, unsigned int* d_in_bit, int bit, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        d_in_bit[idx] = (d_in[idx] >> bit) & 1;
    }
}

// 反转数组
__global__ void reverseArray(unsigned int* input, unsigned int* output, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        output[n - 1 - idx] = input[idx];
    }
}


// 重排数组
__global__ void rearrangeArray(unsigned int* d_in, unsigned int* d_in_pos, unsigned int* d_in_bit,
                                unsigned int* d_prefix_sum, unsigned int* d_suffix_sum,
                                unsigned int* d_tmp, unsigned int* d_tmp_pos, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        if (d_in_bit[idx] == 0) {
            d_tmp[d_prefix_sum[idx]] = d_in[idx];           // exclusive 前缀和，直接用
            d_tmp_pos[d_prefix_sum[idx]] = d_in_pos[idx];
        } else {
            d_tmp[n - 1 - d_suffix_sum[idx]] = d_in[idx];  // exclusive 后缀和
            d_tmp_pos[n - 1 - d_suffix_sum[idx]] = d_in_pos[idx];
        }
    }
}

// 将 block 偏移加到输出
__global__ void add_block_offsets(unsigned int* output, unsigned int* block_offsets, int n) {
    int tid = threadIdx.x;
    int bid = blockIdx.x;
    int gid = bid * BLOCK_SIZE + tid;
    if (gid < n) {
        output[gid] += block_offsets[bid];
    }
}

// 辅助kernel：对输入取反 (bit=0 -> 1, bit=1 -> 0)
__global__ void kernel_inverse(unsigned int* input, unsigned int* output, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        output[idx] = (input[idx] == 0) ? 1 : 0;
    }
}

// Device: Blelloch 扫描（exclusive scan）
__device__ unsigned int blelloch_scan_exclusive(unsigned int* s_data, int tid, int block_size) {
    int offset = 1;
    // Up-sweep
    for (int d = block_size >> 1; d > 0; d >>= 1) {
        if (tid < d) {
            int ai = offset * (2 * tid + 1) - 1;
            int bi = offset * (2 * tid + 2) - 1;
            s_data[bi] += s_data[ai];
        }
        offset *= 2;
        __syncthreads();
    }
    unsigned int total = s_data[block_size - 1];
    if (tid == 0) s_data[block_size - 1] = 0;
    __syncthreads();
    // Down-sweep
    for (int d = 1; d < block_size; d <<= 1) {
        offset >>= 1;
        if (tid < d) {
            int ai = offset * (2 * tid + 1) - 1;
            int bi = offset * (2 * tid + 2) - 1;
            unsigned int t = s_data[ai];
            s_data[ai] = s_data[bi];
            s_data[bi] += t;
        }
        __syncthreads();
    }
    return total;
}

// Block 内 exclusive 扫描，并输出 block 总和
__global__ void blelloch_scan_block(unsigned int* input, unsigned int* output, unsigned int* block_sums, int n, bool exclusive) {
    extern __shared__ unsigned int s_data[];
    int tid = threadIdx.x;
    int bid = blockIdx.x;
    int gid = bid * BLOCK_SIZE + tid;

    if (gid < n) {
        s_data[tid] = input[gid];
    } else {
        s_data[tid] = 0;
    }
    __syncthreads();

    unsigned int block_sum = blelloch_scan_exclusive(s_data, tid, BLOCK_SIZE);

    if (gid < n) {
        if(exclusive){
            output[gid] = s_data[tid];  // exclusive
        }else{
            output[gid] = s_data[tid] + input[gid];  // inclusive
        }
    }

    if (tid == 0) {
        block_sums[bid] = block_sum;
    }
}

// 递归扫描主函数（exclusive）
void exclusivePrefixSum(unsigned int* d_input, unsigned int* d_output, int n) {
    if (n <= 1) {
        if (n == 1) d_output[0] = 0;
        return;
    }

    int num_blocks = (n + BLOCK_SIZE - 1) / BLOCK_SIZE;
    size_t shared_mem = BLOCK_SIZE * sizeof(unsigned int);

    unsigned int *d_block_sums;
    CUDA_CHECK(cudaMalloc(&d_block_sums, num_blocks * sizeof(unsigned int)));

    // Step 1: 每个 block 做 exclusive 扫描
    blelloch_scan_block<<<num_blocks, BLOCK_SIZE, shared_mem>>>(d_input, d_output, d_block_sums, n, true);
    CUDA_CHECK(cudaGetLastError());
    cudaDeviceSynchronize();

    // Step 2: 如果 block 数 > 1，递归扫描 block_sums 得到偏移
    if (num_blocks > 1) {
        unsigned int *d_block_offsets;
        CUDA_CHECK(cudaMalloc(&d_block_offsets, num_blocks * sizeof(unsigned int)));

        exclusivePrefixSum(d_block_sums, d_block_offsets, num_blocks); // 递归 exclusive scan
        cudaDeviceSynchronize();

        // Step 3: 将 block_offsets 加到输出上
        add_block_offsets<<<num_blocks, BLOCK_SIZE>>>(d_output, d_block_offsets, n);
        CUDA_CHECK(cudaGetLastError());
        cudaDeviceSynchronize();

        cudaFree(d_block_offsets);
    }

    // Step 4: 清理
    cudaFree(d_block_sums);
}

// exclusive 后缀和：reverse + exclusive scan + reverse
void exclusiveSuffixSum(unsigned int* d_input, unsigned int* d_output, int n) {
    if (n == 0) return;
    unsigned int *d_rev = nullptr, *d_rev_scan = nullptr;
    CUDA_CHECK(cudaMalloc(&d_rev, n * sizeof(unsigned int)));
    CUDA_CHECK(cudaMalloc(&d_rev_scan, n * sizeof(unsigned int)));

    dim3 block(BLOCK_SIZE), grid((n + BLOCK_SIZE - 1) / BLOCK_SIZE);
    reverseArray<<<grid, block>>>(d_input, d_rev, n);
    cudaDeviceSynchronize();

    exclusivePrefixSum(d_rev, d_rev_scan, n);
    cudaDeviceSynchronize();

    reverseArray<<<grid, block>>>(d_rev_scan, d_output, n);
    cudaDeviceSynchronize();

    cudaFree(d_rev);
    cudaFree(d_rev_scan);
}

// GPU 基数排序主函数
void gpu_radix_sort(unsigned int* h_in, unsigned int* h_in_pos,
                    unsigned int* h_out, unsigned int* h_out_pos, int N) {
    unsigned int *d_in, *d_in_pos, *d_tmp, *d_tmp_pos, *d_in_bit, *d_prefix_sum, *d_suffix_sum, *d_temp_bit;
    CUDA_CHECK(cudaMalloc(&d_in, N * sizeof(unsigned int)));
    CUDA_CHECK(cudaMalloc(&d_in_pos, N * sizeof(unsigned int)));
    CUDA_CHECK(cudaMalloc(&d_tmp, N * sizeof(unsigned int)));
    CUDA_CHECK(cudaMalloc(&d_tmp_pos, N * sizeof(unsigned int)));
    CUDA_CHECK(cudaMalloc(&d_in_bit, N * sizeof(unsigned int)));
    CUDA_CHECK(cudaMalloc(&d_prefix_sum, N * sizeof(unsigned int)));
    CUDA_CHECK(cudaMalloc(&d_suffix_sum, N * sizeof(unsigned int)));
    CUDA_CHECK(cudaMalloc(&d_temp_bit, N * sizeof(unsigned int)));

    CUDA_CHECK(cudaMemcpy(d_in, h_in, N * sizeof(unsigned int), cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_in_pos, h_in_pos, N * sizeof(unsigned int), cudaMemcpyHostToDevice));

    int numBlocks = (N + BLOCK_SIZE - 1) / BLOCK_SIZE;

    for (int bit = 0; bit < MAXBIT; bit++) {
        computeBitValues<<<numBlocks, BLOCK_SIZE>>>(d_in, d_in_bit, bit, N);
        cudaDeviceSynchronize();

        // 前缀和：统计 bit == 0 的数量（exclusive）
        kernel_inverse<<<numBlocks, BLOCK_SIZE>>>(d_in_bit, d_temp_bit, N);
        cudaDeviceSynchronize();
        exclusivePrefixSum(d_temp_bit, d_prefix_sum, N);

        // 后缀和：统计 bit == 1 的数量（exclusive）
        exclusiveSuffixSum(d_in_bit, d_suffix_sum, N);

        // 重排
        rearrangeArray<<<numBlocks, BLOCK_SIZE>>>(d_in, d_in_pos, d_in_bit, d_prefix_sum, d_suffix_sum, d_tmp, d_tmp_pos, N);
        cudaDeviceSynchronize();

        std::swap(d_in, d_tmp);
        std::swap(d_in_pos, d_tmp_pos);
    }

    CUDA_CHECK(cudaMemcpy(h_out, d_in, N * sizeof(unsigned int), cudaMemcpyDeviceToHost));
    CUDA_CHECK(cudaMemcpy(h_out_pos, d_in_pos, N * sizeof(unsigned int), cudaMemcpyDeviceToHost));

    // 释放内存
    cudaFree(d_in); cudaFree(d_in_pos); cudaFree(d_tmp); cudaFree(d_tmp_pos);
    cudaFree(d_in_bit); cudaFree(d_prefix_sum); cudaFree(d_suffix_sum); cudaFree(d_temp_bit);
}

int main() {
    // unsigned int* h_in = new unsigned int[DATASIZE] {5, 3, 7, 2, 9};
    // unsigned int* h_in_pos = new unsigned int[DATASIZE] {0, 1, 2, 3, 4};
    unsigned int* h_in = new unsigned int[DATASIZE];
    unsigned int* h_in_pos = new unsigned int[DATASIZE];
    unsigned int* h_out = new unsigned int[DATASIZE];
    unsigned int* h_out_pos = new unsigned int[DATASIZE];
    unsigned int* cpu_out = new unsigned int[DATASIZE];
    unsigned int* cpu_out_pos = new unsigned int[DATASIZE];

    srand(100);
    for (int i = 0; i < DATASIZE; ++i) {
        h_in[i] = rand() % (DATASIZE + 1);
        h_in_pos[i] = i;
    }

    auto start_cpu = high_resolution_clock::now();
    cpu_radix_sort_with_index(h_in, h_in_pos, cpu_out, cpu_out_pos, DATASIZE);
    auto end_cpu = high_resolution_clock::now();
    auto cpu_time = duration_cast<milliseconds>(end_cpu - start_cpu).count();

    auto start = high_resolution_clock::now();
    gpu_radix_sort(h_in, h_in_pos, h_out, h_out_pos, DATASIZE);
    auto end = high_resolution_clock::now();
    auto gpu_time = duration_cast<milliseconds>(end - start).count();

    print_array(h_in, DATASIZE, "originally value:");
    print_array(cpu_out, DATASIZE, "cpu sorted value:");
    print_array(h_out, DATASIZE, "gpu sorted value:");
    print_array(cpu_out_pos, DATASIZE, "cpu sorted index:");
    print_array(h_out_pos, DATASIZE, "gpu sorted index:");

    printf("Sort value: %s\n", validate(h_out, cpu_out, DATASIZE) ? "CORRECT" : "INCORRECT");
    printf("Sort index: %s\n", validate(h_out_pos, cpu_out_pos, DATASIZE) ? "CORRECT" : "INCORRECT");
    std::cout << "Data size: " << DATASIZE << " elements" << std::endl;
    std::cout << "GPU  time: " << gpu_time << " ms" << std::endl;
    std::cout << "CPU  time: " << cpu_time << " ms" << std::endl;

    delete[] h_in;
    delete[] h_in_pos;
    delete[] h_out;
    delete[] h_out_pos;
    delete[] cpu_out;
    delete[] cpu_out_pos;

    return 0;
}